# Author Zenos
# Create 2025/2/28 上午11:44
import torchvision
from torch import nn
import numpy as np
import torch
import matplotlib.pyplot as plt
from dataset import FontDataset
from utils.utils import custom_collate_fn


def bilinear_kernel(in_channels, out_channels, kernel_size):
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size,
                       kernel_size), dtype='float32')
    weight[range(in_channels), range(out_channels), :, :] = filt
    return torch.from_numpy(weight)


class mynet(nn.Module):
    def __init__(self, num_classes):
        super(mynet, self).__init__()
        pretrained_net = torchvision.models.resnet34(weights=None)
        self.stage = list(pretrained_net.children())[:4]  # 第一段
        # 输入 [1,320,320]
        # 修改第一层的通道数  输出[1, 64, 160, 160]
        self.stage[0] = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        # 将第四层最大池化改成空洞卷积   输出[1, 64, 80, 80]
        self.stage[3] = nn.Conv2d(64, 64, 2, stride=2, padding=1, bias=False, dilation=2)
        # 加入一层批标准化
        batch = nn.BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        self.stage.append(batch)
        # # 转换成Sequential
        self.stage = nn.Sequential(*self.stage)
        # 加入
        self.layer1 = list(pretrained_net.children())[4]  # layer1   输出[64,80,80]
        self.layer2 = list(pretrained_net.children())[5]  # layer2   输出[64,40,40]
        self.layer3 = list(pretrained_net.children())[6]  # layer3   输出[64,20,20]
        self.layer4 = list(pretrained_net.children())[7]  # layer4   输出[64,10,10]

        # 改变通道数
        self.scores1 = nn.Conv2d(512, num_classes, 1)
        self.scores2 = nn.Conv2d(256, num_classes, 1)
        self.scores3 = nn.Conv2d(128, num_classes, 1)
        self.scores4 = nn.Conv2d(64, num_classes, 1)

        # 2倍上采样
        self.upsample_2x = nn.ConvTranspose2d(
            num_classes, num_classes, 4, 2, 1, bias=False)
        self.upsample_2x.weight.data = bilinear_kernel(
            num_classes, num_classes, 4)

    # def forward(self, x, c):
    #     outputs = []
    #     x = self.stage(x)         # [1, 64, 80, 80]
    #     outputs.append(x)
    #     x = self.layer1(x)        # [1, 64, 80, 80]
    #     s1 = x
    #     x = self.layer2(x)        # [1, 128, 40, 40]
    #     s2 = x
    #     x = self.layer3(x)        # [1, 256, 20, 20]
    #     s3 = x
    #     x = self.layer4(x)        # [1, 512, 10, 10]
    #     outputs.append(x)
    #     x = self.scores1(x)       # [1, 35, 10, 10]
    #     # 改变每层下采样通道数
    #     s3 = self.scores2(s3)     # [1, 35, 20, 20]
    #     s2 = self.scores3(s2)     # [1, 35, 40, 40]
    #     s1 = self.scores4(s1)     # [1, 35, 80, 80]
    #     x = self.upsample_2x(x)   # [1, 35, 20, 20]
    #     x = x + s3
    #     x = self.upsample_2x(x)   # [1, 35, 40, 40]
    #     x = x + s2
    #     x = self.upsample_2x(x)   # [1, 35, 80, 80]
    #     outputs.append(x)
    #     x = x + s1
    #     x = self.upsample_2x(x)   # [1, 35, 160, 160]
    #     x = self.upsample_2x(x)   # [1, 35, 320, 320]
    #     outputs.append(x)
    #     return x, outputs

    def forward(self, x, c):
        x = self.stage(x)         # [1, 64, 80, 80]
        x = self.layer1(x)        # [1, 64, 80, 80]
        s1 = x
        x = self.layer2(x)        # [1, 128, 40, 40]
        s2 = x
        x = self.layer3(x)        # [1, 256, 20, 20]
        s3 = x
        x = self.layer4(x)        # [1, 512, 10, 10]
        x = self.scores1(x)       # [1, 35, 10, 10]
        # 改变每层下采样通道数
        s3 = self.scores2(s3)     # [1, 35, 20, 20]
        s2 = self.scores3(s2)     # [1, 35, 40, 40]
        s1 = self.scores4(s1)     # [1, 35, 80, 80]
        x = self.upsample_2x(x)   # [1, 35, 20, 20]
        x = x + s3
        x = self.upsample_2x(x)   # [1, 35, 40, 40]
        x = x + s2
        x = self.upsample_2x(x)   # [1, 35, 80, 80]
        x = x + s1
        x = self.upsample_2x(x)   # [1, 35, 160, 160]
        x = self.upsample_2x(x)   # [1, 35, 320, 320]
        return x


if __name__ == "__main__":
    TrainDataset = FontDataset(True, "/Users/zenos/Downloads/CCSSD/FZLBJW2017")
    # 定义数据集迭代器
    train_iter = torch.utils.data.DataLoader(TrainDataset, 1, shuffle=False, drop_last=True,
                                             collate_fn=custom_collate_fn)

    # 创建网络实例
    net = mynet(35)

    # 从数据集中取一批数据
    for batch in train_iter:
        input_tensor, _, _ = batch  # 假设你的 FontDataset 返回 (image, label)

        # 确保输入是 (batch_size=1, channel=1, height, width)
        if input_tensor.dim() == 3:
            input_tensor = input_tensor.unsqueeze(0)

        # 输入网络
        _, feature_maps = net(input_tensor, "c")

        # 可视化每一层的特征图
        for i, feature_map in enumerate(feature_maps):
            plt.figure(figsize=(15, 15))
            num_channels = feature_map.shape[1]
            for j in range(min(num_channels, 64)):
                plt.subplot(8, 8, j + 1)
                plt.imshow(feature_map[0, j].detach().cpu().numpy(), cmap='jet')
                plt.axis('off')
            plt.suptitle(f'Feature Map after Layer {i + 1}')
            plt.show()

        break  # 只取一张图片，避免循环
